from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph,START,END,add_messages

from typing import TypedDict,List,Annotated
import operator
import os


# 设置环境变量
os.environ["LANGCHAIN_TRACING_V2"] = "true"
os.environ["LANGCHAIN_API_KEY"] = "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"
os.environ["LANGCHAIN_PROJECT"] = "chat1test"
#os.environ["LANGCHAIN_ENDPOINT"]= "https://api.smith.langchain.com"

# 验证环境变量是否设置成功
print(os.getenv("LANGCHAIN_TRACING_V2"))
print(os.getenv("LANGCHAIN_API_KEY"))
print(os.getenv("LANGCHAIN_PROJECT"))

class State(TypedDict):
    messages:Annotated[List[str],operator.add]

builder = StateGraph(State)

api_key = "sk-6S0PtpNia71gjcfwSsDPsJ9mGqsVPr2XRQzAx1dHbJS7RW4t"
api_base="https://chatapi.littlewheat.com/v1"

llm = ChatOpenAI(model="gpt-4o",api_key=api_key,base_url=api_base,temperature=0)

def chatbot(state:State):
    print(state)
    return {"messages":[llm.invoke(state["messages"])]}

builder.add_node("chatbot",chatbot)
builder.add_edge(START,"chatbot")
builder.add_edge("chatbot",END)

graph = builder.compile()

def stream_graph_updates(user_input:str):
    for event in graph.stream({"messages":[("user",user_input)]}):
        for value in event.values():
            print("模型回复:", value["messages"][-1].content)

while True:
    try:
        user_input = input("用户输入：")
        if user_input.lower() in ["退出"]:
            print("下次再见")
            break

        stream_graph_updates(user_input)
    except:
        break
